package cdu.sl.aichatdemo.service.impl;

import cdu.sl.aichatdemo.mapper.CustomChatMemoryMapper;
import cdu.sl.aichatdemo.pojo.entity.CustomChatMemory;
import cdu.sl.aichatdemo.service.TitleGenerateService;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Description 自定义聊天记忆实现
 * @Author sunlin5@asiainfo.com
 * @Date 2025/8/28
 * @Version 1.0
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class CustomChatMemoryImpl implements ChatMemory {

    private final CustomChatMemoryMapper customChatMemoryMapper;

    private final TitleGenerateService titleGenerateService;

    // 临时写死的用户ID，后续可以从上下文获取
    private static final Long DEFAULT_USER_ID = 1L;

    @Override
    public void add(String conversationId, Message message) {
        log.info("新增消息，conversationId: {}, messageType: {}", conversationId, message.getClass().getSimpleName());

        // 获取当前对话中的消息数量，用于设置orderIndex
        QueryWrapper<CustomChatMemory> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("conversation_id", conversationId);
        Long count = customChatMemoryMapper.selectCount(queryWrapper);

        CustomChatMemory chatMemory = new CustomChatMemory();
        chatMemory.setConversationId(conversationId);
        chatMemory.setUserId(DEFAULT_USER_ID);
        chatMemory.setContent(message.getText());
        chatMemory.setOrderIndex(count.intValue() + 1);
        chatMemory.setCreatedAt(LocalDateTime.now());
        chatMemory.setUpdatedAt(LocalDateTime.now());

        // 根据消息类型设置type字段
        if (message instanceof UserMessage) {
            chatMemory.setType("USER");
        } else if (message instanceof AssistantMessage) {
            chatMemory.setType("ASSISTANT");
        } else {
            chatMemory.setType("SYSTEM");
        }

        customChatMemoryMapper.insert(chatMemory);
        // 生成标题
        // 如果考虑优化 将对话列表和对话内容拆成两张表 对话列表只维护标题 id 内容数量 是否生成过标题等统计性信息
        // 对话内容表存储具体的对话内容
        titleGenerateService.generateTitle(conversationId);
        log.info("消息保存成功，ID: {}", chatMemory.getId());
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        log.info("新增消息列表，conversationId: {}, 消息数量: {}", conversationId, messages.size());

        // 获取当前对话中的消息数量
        QueryWrapper<CustomChatMemory> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("conversation_id", conversationId);
        Long count = customChatMemoryMapper.selectCount(queryWrapper);

        AtomicInteger orderIndex = new AtomicInteger(count.intValue() + 1);

        for (Message message : messages) {
            CustomChatMemory chatMemory = new CustomChatMemory();
            chatMemory.setConversationId(conversationId);
            chatMemory.setUserId(DEFAULT_USER_ID);
            chatMemory.setContent(message.getText());
            chatMemory.setOrderIndex(orderIndex.getAndIncrement());
            chatMemory.setCreatedAt(LocalDateTime.now());
            chatMemory.setUpdatedAt(LocalDateTime.now());

            // 根据消息类型设置type字段
            if (message instanceof UserMessage) {
                chatMemory.setType("USER");
            } else if (message instanceof AssistantMessage) {
                chatMemory.setType("ASSISTANT");
            } else {
                chatMemory.setType("SYSTEM");
            }

            customChatMemoryMapper.insert(chatMemory);
            titleGenerateService.generateTitle(conversationId);
        }
        log.info("消息列表保存成功");
    }

    @Override
    public void clear(String conversationId) {
        log.info("清空消息，conversationId: {}", conversationId);

        QueryWrapper<CustomChatMemory> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("conversation_id", conversationId);


        CustomChatMemory updateEntity = new CustomChatMemory();
        updateEntity.setUpdatedAt(LocalDateTime.now());

        customChatMemoryMapper.update(updateEntity, queryWrapper);
        log.info("对话记录清空成功");
    }

    @Override
    public List<Message> get(String conversationId) {
        log.info("根据对话id获取消息，conversationId: {}", conversationId);

        QueryWrapper<CustomChatMemory> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("conversation_id", conversationId)
                .orderByAsc("order_index");

        List<CustomChatMemory> chatMemories = customChatMemoryMapper.selectList(queryWrapper);

        List<Message> messages = new ArrayList<>();
        for (CustomChatMemory chatMemory : chatMemories) {
            Message message;
            if ("USER".equals(chatMemory.getType())) {
                message = new UserMessage(chatMemory.getContent());
            } else if ("ASSISTANT".equals(chatMemory.getType())) {
                message = new AssistantMessage(chatMemory.getContent());
            } else {
                // 对于SYSTEM类型的消息，使用UserMessage作为默认
                message = new UserMessage(chatMemory.getContent());
            }
            messages.add(message);
        }

        log.info("获取到消息数量: {}", messages.size());
        return messages;
    }
}
